import pickle as pkl
import time
from datetime import date, timedelta, datetime
from io import StringIO

import numpy as np
import pandas as pd
import streamlit as st

from enigma import CHARS, Enigma
from gen_codebook import generate_codebook


def parse_codebook(codebook_file):
    # TODO：格式检查
    configure = {}
    codebook = pd.read_csv(codebook_file)
    for line in codebook.values:
        config = {}
        config['wheel_num'] = [x for x in str(line[1])]
        config['carry_state'] = [x for x in line[2]]
        config['init_state'] = [x for x in line[3]]
        config['single_table'] = [[y for y in x] for x in line[4].split(' ')]
        configure[str(line[0])] = config

    return configure


st.title('欢迎使用Enigmar')

if 'config' not in st.session_state:
    st.session_state.config = None

with st.sidebar:
    st.header('控制台')

    reset = st.button('重置全部设置')
    if reset:
        st.experimental_rerun()

    st.markdown('---')

    st.subheader('Enigma型号')

    enigma_type = st.selectbox(
        '',
        ['3轮', '4轮']
    )

    if_single_tabel = st.checkbox('单表连接板', value=True)

    # TODO: 增加型号
    if enigma_type == '3轮':
        num_wheels = 3
        wheels_file = 'wheels.pkl'
    elif enigma_type == '4轮':
        num_wheels = 4
        wheels_file = 'wheels.pkl'

    with open(wheels_file, 'rb') as fw:
        all_wheels = pkl.load(fw)
    
    st.markdown('---')

    st.subheader('密钥本设置')

    col_u, col_v = st.columns(2)
    with col_u:
        st.markdown('##### 生成密钥本')
        start_date = st.date_input('从', value=date.today())
        end_date = st.date_input('到', value=date.today() + timedelta(days=1))
        correct_date = False
        if start_date.__gt__(end_date):
            st.error('生成开始日期不能晚于结束日期')
        else:
            correct_date = True
        start_gen = st.button(
            '生成',
            disabled=not correct_date,
            help='生成指定时间段内的密码本，以天为单位'
            )
        if start_gen:
            start_time = int(time.mktime(start_date.timetuple()))
            end_time = int(time.mktime(end_date.timetuple())) + 86400
            codebook = generate_codebook(num_wheels, start_time, end_time)
            codebook_data = 'date,wheel_num,carry_state,init_state,single_table\n' + ''.join(codebook)
            st.write('生成完毕！')
            download_codebook = st.download_button(
                '导出',
                data=codebook_data,
                file_name='codebook.csv',
                disabled=not start_gen
                )
    with col_v:
        st.markdown('##### 导入密钥本')
        codebook_upload = st.file_uploader(
            '上传密钥本',
            help='请上传需导入的密钥本，密钥本格式样本可从左侧密钥本生成功能导出'
            )
        if codebook_upload:
            try:
                configure = parse_codebook(codebook_upload)
            except:
                st.error('密码本格式错误，请重新上传。请使用密钥本生成功能，建议不要手动修改密钥本')
            config_date = st.date_input(
                '密钥日期',
                min_value=datetime.strptime(sorted(configure.keys())[0], '%Y%m%d'),
                max_value=datetime.strptime(sorted(configure.keys())[-1], '%Y%m%d'),
            )
        start_parse = st.button('导入', disabled=not codebook_upload)
        # st.write(st.session_state)
        config = None
        if start_parse:
            try:
                for k in st.session_state:
                    del st.session_state[k]
                config = configure[config_date.strftime('%Y%m%d')]
            except KeyError as e:
                st.error('未查询到对应日期，请重新输入')
            st.session_state.config = config

    st.markdown('---')

    st.subheader('转子选择')

    selected_wheels = []
    all_wheels_for_show = np.array(list(range(len(all_wheels))))
    all_wheels_for_show += 1
    if_set_error = False
    selected_wheels_origin = st.multiselect(
        '',
        all_wheels_for_show,
        default=[int(x) for x in st.session_state.config['wheel_num']] if st.session_state.config else None,
        key='wheel_selected'
    )
    if len(selected_wheels_origin) > num_wheels:
        st.error(f'您选择的转子过多，请选择{num_wheels}个转子')
        if_set_error=True
    elif len(selected_wheels_origin) < num_wheels:
        st.error(f'您选择的转子过少，请选择{num_wheels}个转子')
        if_set_error=True
    selected_wheels_idx = list(map(lambda x: x - 1, selected_wheels_origin))
    for idx in selected_wheels_idx:
        selected_wheels.append(all_wheels[idx])

    st.markdown('---')

    st.subheader('设置进位点')

    wheels_carry_points = []
    cols_3 = st.columns(num_wheels)
    col_i = 0
    for col in cols_3:
        with col:
            container = st.container()
            origin_idx = CHARS
            show_idx = list(map(chr, origin_idx))
            w_setter = st.selectbox(
                f'第{col_i + 1}转子',
                show_idx,
                index=show_idx.index(st.session_state.config['carry_state'][col_i]) if st.session_state.config else 0,
                key=f'wheel_carry_setter_{col_i}'
            )
            wheels_carry_points.append(ord(w_setter))
            col_i += 1

    st.markdown('---')

    st.subheader('设置初始位置')

    selected_wheels_init = []
    cols_2 = st.columns(num_wheels)
    col_i = 0
    for col in cols_2:
        with col:
            container = st.container()
            origin_idx = CHARS
            show_idx = list(map(chr, origin_idx))
            w_setter = st.selectbox(
                f'第{col_i + 1}转子',
                show_idx,
                index=show_idx.index(st.session_state.config['init_state'][col_i]) if st.session_state.config else 0,
                key=f'wheel_setter_{col_i}')
            selected_wheels_init.append(ord(w_setter))
            col_i += 1


    if if_single_tabel:
        st.markdown('---')

        st.subheader('设置连接板')

        col_m, col_n = st.columns(2)
        num_mapping = 0
        single_table = {}
        char_for_show = list(map(chr, CHARS))
        if st.session_state.config:
            # st.write(list(zip(*st.session_state.config['single_table'])))
            default_singletable_a, default_singletable_b = list(zip(*st.session_state.config['single_table']))
        else:
            default_singletable_a, default_singletable_b = (None, None)
        for c in CHARS:
            single_table[chr(c)] = chr(c)
        with col_m:
            single_table_selected_a = st.multiselect(
                '',
                char_for_show,
                default=default_singletable_a,
                help='两侧对应位置的字母将被交换，同一字母不能在两侧同时被选择（避免多重交换）',
                key='single_table_a'
            )
        with col_n:
            single_table_selected_b = st.multiselect(
                '',
                char_for_show,
                default=default_singletable_b,
                help='两侧对应位置的字母将被交换，同一字母不能在两侧同时被选择（避免多重交换）',
                key='single_table_b'
            )
        
        if_singletabel_error = False
        error_c = set()
        for c in single_table_selected_a:
            if c in single_table_selected_b:
                if_singletabel_error = True
                error_c.add(c)
        for c in single_table_selected_b:
            if c in single_table_selected_a:
                if_singletabel_error = True
                error_c.add(c)
        if if_singletabel_error:
            for c in error_c:
                st.error(f'字符{c}在两侧同时被选择，请更改')
            if_set_error = True

        single_table_selected = list(zip(single_table_selected_a, single_table_selected_b))
        # st.write(single_table_selected)
        if single_table_selected:
            single_table_selected = np.array(single_table_selected).flatten()
        # st.write(single_table_selected)
        char_a = []
        char_b = []
        for n in range(len(single_table_selected)):
            if n % 2 == 0:
                char_a.append(single_table_selected[n])
            else:
                char_b.append(single_table_selected[n])
        for i in range(len(char_a)):
            single_table[char_a[i]] = char_b[i]
            single_table[char_b[i]] = char_a[i]

output_text = ''
if not if_set_error:
    if if_single_tabel:
        enigma = Enigma(selected_wheels, selected_wheels_init, wheels_carry_points, single_table)
    else:
        enigma = Enigma(selected_wheels, selected_wheels_init, wheels_carry_points)
    with st.form('input_text'):
        input_text = st.text_area(
            '',
            height=200,
            max_chars=1024,
            placeholder='只能输入字符集内的字符及空白字符（空格，回车，TAB），长度限制：1024')
        submitted = st.form_submit_button('提交')
    if submitted:
        bar = st.progress(0.)
        try:
            if if_single_tabel:
                enigma = Enigma(selected_wheels, selected_wheels_init, wheels_carry_points, single_table)
            else:
                enigma = Enigma(selected_wheels, selected_wheels_init, wheels_carry_points)
            output_text = enigma.encode(input_text, bar)
        except AssertionError:
            st.error(f'您的输入有误，只能输入字符集内的字符。当前字符集：{list(map(chr, CHARS))}')

    st.write('输出')
    st.code(output_text)
else:
    st.info('请确保设置正确以开始编码')
